EKS オーケストレータを使った最低限の HyperPod クラスターを作成してみる

たかくに
2024.12.29
こんにちは！ AWS 事業本部コンサルティング部のたかくに（@takakuni_）です。
みなさん SageMaker HyperPod 触っていますでしょうか。
SageMaker HyperPod の強みと言えば、オーケスとレーターに EKS が使えます。
Slurm オーケストレータに関しては AWS ParallelCluster や AWS Parallel Computing Service もサポートしていますが、 EKS オーケストレータは SageMaker HyperPod のみサポートしている状況です。
re:Invent 2024 の A close look at how Amazon built the Nova FMs using SageMaker HyperPod (AIM379) では Amazon Nova を SageMaker HyperPod on EKS で構築した話が紹介されています。
https://youtu.be/jXcWKJE1Xts?feature=shared
そんな、 SageMaker HyperPod on EKS を作るにおいて最低限必要なリソースを押さえていきたいと思います。
 構成今回は最低限にこだわるため、以下の構成で作成します。
EKS からみると ENI しか見えない self-managed なノードで管理している不思議な感覚なのですが、このようなイメージになります。
 やってみる今回はいつもどおり再現性があるようにしたいため、 HashiCorp Terraform を利用しました。作成したコードは以下に格納せれています。
https://github.com/takakuni-classmethod/genai-blog/tree/main/sagemaker_hyperpod_101_eks
 IAMSageMaker HyperPod と EKS クラスターの IAM について触れます。
 SageMaker HyperPod 側まずは SageMaker HyperPod 側ですが Slurm オーケストレータと権限が異なります。
マネージドポリシーの AmazonSageMakerClusterInstanceRolePolicy と以下のカスタムポリシーが必要です。
iam_hyperpod.json
{
	"Version": "2012-10-17",
	"Statement": [
		{
			"Effect": "Allow",
			"Action": [
				"ec2:AssignPrivateIpAddresses",
				"ec2:CreateNetworkInterface",
				"ec2:CreateNetworkInterfacePermission",
				"ec2:DeleteNetworkInterface",
				"ec2:DeleteNetworkInterfacePermission",
				"ec2:DescribeNetworkInterfaces",
				"ec2:DescribeVpcs",
				"ec2:DescribeDhcpOptions",
				"ec2:DescribeSubnets",
				"ec2:DescribeSecurityGroups",
				"ec2:DetachNetworkInterface",
				"ec2:ModifyNetworkInterfaceAttribute",
				"ec2:UnassignPrivateIpAddresses",
				"ecr:BatchGetImage",
				"ecr:GetAuthorizationToken",
				"ecr:GetDownloadUrlForLayer",
				"eks-auth:AssumeRoleForPodIdentity"
			],
			"Resource": "*"
		},
		{
			"Effect": "Allow",
			"Action": ["ec2:CreateTags"],
			"Resource": ["arn:aws:ec2:*:*:network-interface/*"]
		}
	]
}
eks-auth:AssumeRoleForPodIdentity は Pod Identity を利用する場合に必要な権限なのでケースバイケースで付け替えてあげてください。
https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-hyperpod-prerequisites-iam.html#sagemaker-hyperpod-prerequisites-iam-role-for-hyperpod
Terraform では以下の部分が該当するコードです。
hyperpod.tf
###################################################
# IAM Role for SageMaker HyperPod Cluster
###################################################
resource "aws_iam_role" "hyperpod" {
  name = "${local.prefix}-hyperpod-role"
  assume_role_policy = jsonencode({
    Version = "2012-10-17",
    Statement = [
      {
        Effect = "Allow",
        Principal = {
          Service = "sagemaker.amazonaws.com"
        },
        Action = "sts:AssumeRole"
      }
    ]
  })
}

resource "aws_iam_role_policy_attachment" "hyperpod_managed" {
  role       = aws_iam_role.hyperpod.name
  policy_arn = "arn:aws:iam::aws:policy/AmazonSageMakerClusterInstanceRolePolicy"
}

data "aws_iam_policy_document" "hyperpod_vpc" {
  statement {
    effect = "Allow"
    actions = [
      "ec2:AssignPrivateIpAddresses",
      "ec2:CreateNetworkInterface",
      "ec2:CreateNetworkInterfacePermission",
      "ec2:DeleteNetworkInterface",
      "ec2:DeleteNetworkInterfacePermission",
      "ec2:DescribeNetworkInterfaces",
      "ec2:DescribeVpcs",
      "ec2:DescribeDhcpOptions",
      "ec2:DescribeSubnets",
      "ec2:DescribeSecurityGroups",
      "ec2:DetachNetworkInterface",
      "ec2:ModifyNetworkInterfaceAttribute",
      "ec2:UnassignPrivateIpAddresses",
      "ecr:BatchGetImage",
      "ecr:GetAuthorizationToken",
      "ecr:GetDownloadUrlForLayer",
      "eks-auth:AssumeRoleForPodIdentity"
    ]
    resources = ["*"]
  }

  statement {
    effect = "Allow"
    actions = [
      "ec2:CreateTags",
    ]
    resources = ["arn:aws:ec2:*:*:network-interface/*"]
  }
}

resource "aws_iam_policy" "hyperpod_vpc" {
  name        = "${local.prefix}-vpc-policy"
  description = "IAM policy for SageMaker HyperPod VPC"
  policy      = data.aws_iam_policy_document.hyperpod_vpc.json
}

resource "aws_iam_role_policy_attachment" "hyperpod_vpc" {
  role       = aws_iam_role.hyperpod.name
  policy_arn = aws_iam_policy.hyperpod_vpc.arn
}

data "aws_iam_policy_document" "hyperpod_lifecycle" {
  statement {
    effect = "Allow"
    actions = [
      "s3:ListBucket",
      "s3:GetObject"
    ]
    resources = [
      aws_s3_bucket.life_cycle_scripts.arn,
      "${aws_s3_bucket.life_cycle_scripts.arn}/*"
    ]
  }
}

resource "aws_iam_policy" "hyperpod_lifecycle" {
  name        = "${local.prefix}-lifecycle-policy"
  description = "IAM policy for SageMaker HyperPod Lifecycle Scripts"
  policy      = data.aws_iam_policy_document.hyperpod_lifecycle.json
}

resource "aws_iam_role_policy_attachment" "hyperpod_lifecycle" {
  role       = aws_iam_role.hyperpod.name
  policy_arn = aws_iam_policy.hyperpod_lifecycle.arn
}
hyperpod_lifecycle という、説明していないリソースが含まれていますね。
これは HyperPod インスタンスグループで利用するライフサイクルスクリプト用の S3 バケットが sagemaker- から始まらない場合に必要な権限になります。詳しくは以下をご覧ください。
https://dev.classmethod.jp/articles/sagemaker-hyperpod-execution-role-s3-bucket-permissions/
If you create an IAM role for SageMaker HyperPod attaching only the managed AmazonSageMakerClusterInstanceRolePolicy, your cluster has access to Amazon S3 buckets with the specific prefix sagemaker-.
https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-hyperpod-eks-operate-cli-command-create-cluster.html
Make sure that the S3 bucket path starts with s3://sagemaker-. The IAM role for SageMaker HyperPod has the managed AmazonSageMakerClusterInstanceRolePolicy attached, which allows access to S3 buckets with the specific prefix sagemaker-.
https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_ClusterLifeCycleConfig.html#sagemaker-Type-ClusterLifeCycleConfig-SourceS3Uri
 EKS 側EKS 側は EKS クラスターをホストするために必要な IAM ポリシーのみで十分です。
特段こだわりがなければ AmazonEKSClusterPolicy を利用しましょう。
https://docs.aws.amazon.com/eks/latest/userguide/cluster-iam-role.html
SageMaker HyperPod on EKS の世界ではノードの管理は HyperPod 側（インスタンスグループの管理）になるため、ノードグループのような概念はないです。
そのため、ノードグループにいつも付与していたワーカーノード用 IAM ロールの作成は不要になります。（HyperPod のインスタンスグループに紐づけた実行ロールが担います。）
eks.tf
###################################################
# IAM Role for EKS Cluster
###################################################
resource "aws_iam_role" "eks" {
  name = "${local.prefix}-eks-role"
  assume_role_policy = jsonencode({
    Version = "2012-10-17",
    Statement = [
      {
        Effect = "Allow",
        Principal = {
          Service = "eks.amazonaws.com"
        },
        Action = "sts:AssumeRole"
      }
    ]
  })
}

resource "aws_iam_role_policy_attachment" "eks_cluster" {
  role       = aws_iam_role.eks.name
  policy_arn = "arn:aws:iam::aws:policy/AmazonEKSClusterPolicy"
}
 セキュリティグループ SageMaker HyperPodHyperPod 側ですがマウントするファイルシステムや EFA の利用有無によって異なります。
どちらも利用しない場合は、各ノードへのメンテナンス用の接続は SSM を介して行うため、インバウンドは不要。資材をダウンロードするためのアウトバウンドが解放されていれば OK です。
hyperpod.tf
###################################################
# Security Group for SageMaker HyperPod Cluster
###################################################
resource "aws_security_group" "hyperpod" {
  name        = "${local.prefix}-hyperpod-sg"
  vpc_id      = module.vpc.vpc_id
  description = "${local.prefix}-hyperpod-sg"

  tags = {
    Name = "${local.prefix}-hyperpod-sg"
  }
}

resource "aws_vpc_security_group_egress_rule" "hyperpod_allow_all_traffic_ipv4" {
  security_group_id = aws_security_group.hyperpod.id
  cidr_ipv4         = "0.0.0.0/0"
  ip_protocol       = "-1"
}
なお、 HyperPod on EKS でホストするノードは、プライベートサブネットな必要があります。そのため、パブリックサブネットに配置し、パブリック IP 直繋ぎのようなことはできません。
The type of the subnet in your VPC must be private for HyperPod clusters.
https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-hyperpod-eks-prerequisites.html
 EKS続いて EKS コントロールプレーン側のクラスターセキュリティグループです。こちらも特段 SageMaker HyperPod と通信を行うわけではないため、アウトバウンドが空いていれば OK です。
eks.tf
###################################################
# Security Group for EKS Cluster
###################################################
resource "aws_security_group" "eks" {
  name        = "${local.prefix}-eks-sg"
  vpc_id      = module.vpc.vpc_id
  description = "${local.prefix}-eks-sg"

  tags = {
    Name = "${local.prefix}-eks-sg"
  }
}

resource "aws_vpc_security_group_egress_rule" "eks_allow_all_traffic_ipv4" {
  security_group_id = aws_security_group.eks.id
  cidr_ipv4         = "0.0.0.0/0"
  ip_protocol       = "-1"
}
 依存関係をインストールHyperPod で利用するネームスペースやコンポーネントを EKS クラスターにインストールする必要があります。
AWS から以下が用意されているため基本的にはこちらを利用しましょう。
https://github.com/aws/sagemaker-hyperpod-cli/tree/main/helm_chart
私の場合は Helm プロバイダーを利用しました。
ただし、Helm プロバイダーはデフォルトではインストールを同期的に行うため循環が起きています。（HyperPod のためのインストールだが、ノードの管理は HyperPod 側で管理されている。）
そこで、インストールを非同期で行うよう、 wait = false に設定を行いました。
providers.tf (抜粋)
terraform {
  required_providers {
    helm = {
      source  = "hashicorp/helm"
      version = "2.17.0"
    }
  }
}

provider "helm" {
  kubernetes {
    host                   = aws_eks_cluster.this.endpoint
    cluster_ca_certificate = base64decode(aws_eks_cluster.this.certificate_authority[0].data)
    token                  = data.aws_eks_cluster_auth.this.token
  }
}
helm.tf
resource "helm_release" "hyperpod_dependencies" {
  name              = "hyperpod-dependencies"
  chart             = "./sagemaker-hyperpod-cli/helm_chart/HyperPodHelmChart"
  dependency_update = true
+  wait              = false

  depends_on = [
    aws_eks_cluster.this,
    aws_eks_access_entry.this
  ]
}
https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-hyperpod-eks-install-packages-using-helm-chart.html
 on_create.shon_create.sh は Slurm オーケスとレーターに比べて非常にシンプルな作りになっています。必要に応じてデフォルトの on_create.sh に追記しましょう。
https://github.com/aws-samples/awsome-distributed-training/blob/main/1.architectures/7.sagemaker-hyperpod-eks/LifecycleScripts/base-config/on_create.sh
s3.tf
###################################################
# Lifecycle Script Bucket for SageMaker HyperPod Cluster
###################################################
resource "aws_s3_bucket" "life_cycle_scripts" {
  bucket = "${local.prefix}-lifecycle-${local.account_id}"
}

resource "aws_s3_bucket_public_access_block" "life_cycle_scripts" {
  bucket                  = aws_s3_bucket.life_cycle_scripts.bucket
  block_public_acls       = true
  block_public_policy     = true
  ignore_public_acls      = true
  restrict_public_buckets = true
}

resource "aws_s3_bucket_ownership_controls" "life_cycle_scripts" {
  bucket = aws_s3_bucket.life_cycle_scripts.bucket
  rule {
    object_ownership = "BucketOwnerPreferred"
  }
}

###################################################
# Lifecycle Script Objects for SageMaker HyperPod Cluster
###################################################
resource "aws_s3_object" "life_cycle_scripts" {
  for_each = fileset("./config/", "**")
  bucket   = aws_s3_bucket.life_cycle_scripts.bucket
  key      = "config/${each.value}"
  source   = "./config/${each.value}"
  etag     = filemd5("./config/${each.value}")
}
 EKS クラスターとうとう、EKS クラスターまで来ました。
いくつか制約事項があります。
 EKS バージョンSageMaker HyperPod でサポートされる EKS クラスターは 2024 年 12 月時点では 1.28, 1.29, 1.30 までです。
SageMaker HyperPod supports Kubernetes versions 1.28, 1.29, and 1.30.
作成する時期によってこの辺りは異なるため、よく確認して EKS クラスターを作成しましょう。
 コントロールプレーンへのアクセスコントロールプレーンのアクセスは API または API_AND_CONFIG_MAP のどちらかを選べます。
API を選べば HyperPod 側で自動できにアクセスエントリが作られます。このエントリは IaC での書き換えや事前登録不可なのでご注意ください。
The authentication mode of an Amazon EKS cluster supported by SageMaker HyperPod are API and API_AND_CONFIG_MAP.
 CNI PluginEKS クラスターにアドオンで CNI Plugin が必要です。v1.18.3 以上をアドオンとしてインストールしておきましょう。
SageMaker HyperPod requires the Amazon VPC Container Network Interface (CNI) plug-in version 1.18.3 or later.
また、その他の add-ons も利用できるため非常に便利ですね。
You can continue using the various add-ons provided by Amazon EKS such as Kube-proxy, CoreDNS, the Amazon VPC Container Network Interface (CNI) plugin, Amazon EKS pod identity, the GuardDuty agent, the Amazon FSx Container Storage Interface (CSI) driver, the Mountpoint for Amazon S3 CSI driver, the AWS Distro for OpenTelemetry, and the CloudWatch Observability agent.
 ストレージについてHyperPod のノードでホストした Pod は直接 EBS ボリュームをマウントできないため、インスタンスグループで追加の EBS ボリュームを設定し Pod 側の local path でマウントしてあげる必要があります。
Considerations for configuring SageMaker HyperPod clusters with Amazon EKS
You can't mount additional EBS volumes directly to Pods running on HyperPod cluster nodes. Instead, you need to utilize InstanceStorageConfigs to provision and mount additional EBS volumes to the HyperPod nodes. It's important to note that you can only attach additional EBS volumes to new instance groups while creating or updating a HyperPod cluster. Once you have configured instance groups with these additional EBS volumes, in your Amazon EKS Pod configuration file, you'll need to set the local path to /opt/sagemaker to properly mount the volumes to your Amazon EKS Pods.
You can deploy the Amazon EBS CSI (Container Storage Interface) controller on HyperPod nodes. However, the Amazon EBS CSI node DaemonSet, which facilitates the mounting and unmounting of EBS volumes, can only run on non-HyperPod instances. If you use instance-type labels for defining scheduling constraints, ensure that you use the SageMaker AI ML instance types prefixed with ml.. For example, for P5 instances, use ml.p5.48xlarge instead of p5.48xlarge.
その他制約事項含め、以下にまとまっているのでご覧ください。
https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-hyperpod-eks-prerequisites.html
eks.tf
###################################################
# EKS Cluster
###################################################
resource "aws_eks_cluster" "this" {
  name = "${local.prefix}-eks-cluster"

  access_config {
    authentication_mode = "API"
  }

  role_arn = aws_iam_role.eks.arn
  version  = "1.30"

  vpc_config {
    security_group_ids = [aws_security_group.eks.id]
    subnet_ids         = module.vpc.private_subnets
  }

  enabled_cluster_log_types = [
    "api",
    "audit",
    "authenticator",
    "controllerManager",
    "scheduler"
  ]

  depends_on = [
    aws_iam_role_policy_attachment.eks_cluster,
  ]
}

###################################################
# EKS Cluster Addons
###################################################
resource "aws_eks_addon" "vpc_cni" {
  cluster_name                = aws_eks_cluster.this.name
  addon_name                  = "vpc-cni"
  addon_version               = "v1.19.2-eksbuild.1"
  resolve_conflicts_on_create = "OVERWRITE"
}

###################################################
# EKS Cluster Access Entries
###################################################
data "aws_iam_session_context" "this" {
  arn = data.aws_caller_identity.self.arn
}

resource "aws_eks_access_entry" "this" {
  cluster_name  = aws_eks_cluster.this.name
  principal_arn = data.aws_iam_session_context.this.issuer_arn
  type          = "STANDARD"
}

resource "aws_eks_access_policy_association" "this" {
  cluster_name  = aws_eks_cluster.this.name
  principal_arn = aws_eks_access_entry.this.principal_arn
  policy_arn    = "arn:aws:eks::aws:cluster-access-policy/AmazonEKSClusterAdminPolicy"

  access_scope {
    type = "cluster"
  }
}
 HyperPod クラスター最後に HyperPod クラスターです。
AWS Provider でサポートしていないため、 AWS Cloud Control Provider を利用します。 EKS クラスター側で制約事項は説明したため他に説明することはないのですが、 source_s3_uri のパスは https:// から始まっても良いみたいです。
SourceS3Uri

Pattern: ^(https|s3)://([^/]+)/?(.*)$
https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_ClusterLifeCycleConfig.html#sagemaker-Type-ClusterLifeCycleConfig-SourceS3Uri
hyperpod.tf
###################################################
# SageMaker HyperPod Cluster
###################################################
resource "awscc_sagemaker_cluster" "this" {
  cluster_name = "${local.prefix}-hyperpod-cluster"

  orchestrator = {
    eks = {
      cluster_arn = aws_eks_cluster.this.arn
    }
  }

  vpc_config = {
    security_group_ids = [aws_security_group.hyperpod.id]
    subnets            = [module.vpc.private_subnets[0]]
  }

  node_recovery = "Automatic" # Automatic | None
  # https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateCluster.html#sagemaker-CreateCluster-request-NodeRecovery

  instance_groups = [
    {
      execution_role      = aws_iam_role.hyperpod.arn
      instance_count      = 2
      instance_group_name = "worker-group"
      instance_type       = "ml.m5.2xlarge"
      life_cycle_config = {
        source_s3_uri = "s3://${aws_s3_bucket.life_cycle_scripts.id}/config/"
        on_create     = "on_create.sh"
      }
      instance_storage_configs = [{
        ebs_volume_config = {
          volume_size_in_gb = 500
        }
      }]
    },
  ]

  depends_on = [
    aws_iam_role_policy_attachment.hyperpod_vpc,
    aws_iam_role_policy_attachment.hyperpod_lifecycle,
    aws_vpc_endpoint.s3
  ]
}
 コンソールを確認できあがったリソースを見てみましょう。
 EKS クラスターSelf-managed なインスタンスが登録されていますね。ノードグループは無しで動いています。
ノードの詳細です。OS は Amazon Linux 2 になっていますね。
Pod では先ほど Helm でインストールしたものも動いていますね。
IAM アクセスエントリには HyperPod のサービスロールと、インスタンスグループで指定した実行ロールが自動的に許可されています。
ß
 SageMaker HyperPod クラスターHyperPod クラスターのインスタンスの状況です。ここは Slurm とあまり変わりがないですね。
設定部分はアドオンのインストール状況や EKS クラスターの詳細が記載されていました。タスクガバナンスはまた別途試してみたいと思います。
 CloudWatch LogsCloudWatch Logs にはインスタンスグループで実行したライフスクリプトのログが記録されます。この辺りを見ながらトラブルシューティングに役立ちそうです。
 まとめ以上、「EKS オーケストレータを使った最低限の HyperPod クラスターを作成してみる」でした。オーケストレータ別にまとめてみました。タスクガバナンスや FSx for Lustre あたりはまだ触れていないため別の機会にチャレンジできればと思います。
AWS 事業本部コンサルティング部のたかくに（@takakuni_）でした！